import pandas as pd
from sklearn.model_selection import train_test_split

# 读取 CSV 文件
input_csv = "/mnt/new/liufenglin/lishiyang/Grounded-Segment-Anything/target_data/data_path_h800_train.csv"  # 替换为你的 CSV 文件路径
df = pd.read_csv(input_csv)

# 确保 'path' 列存在
if 'path' not in df.columns:
    raise ValueError("CSV 文件中没有 'path' 列")

# 随机打乱数据
df = df.sample(frac=1, random_state=42).reset_index(drop=True)

# 按比例分割数据
train_df, temp_df = train_test_split(df, test_size=0.2, random_state=42)
val_df, test_df = train_test_split(temp_df, test_size=0.99, random_state=42)

# 保存为新的 CSV 文件
train_df.to_csv("train.csv", index=False)
val_df.to_csv("val.csv", index=False)
test_df.to_csv("test.csv", index=False)
print(f"训练集长度: {len(train_df)}")
print(f"验证集长度: {len(val_df)}")
print(f"测试集长度: {len(test_df)}")
print("数据分割完成并保存为 train.csv, val.csv, test.csv")